#include <iostream>

#include "UnidirectionLinkedList.hpp"

class Solution
{
public:
    ListNode *partition(ListNode *head, int x)
    {
        if (head == nullptr)
            return nullptr;
        ListNode *lowHead = nullptr, *highHead = nullptr;
        ListNode *lower = nullptr, *higher = nullptr;
        while (head != nullptr)
        {
            if (head->val < x)
            {
                if (lower == nullptr)
                {
                    lower = head;
                    lowHead = lower;
                }
                else
                {
                    lower->next = head;
                    lower = lower->next;
                }
            }
            else
            {
                if (higher == nullptr)
                {
                    higher = head;
                    highHead = higher;
                }
                else
                {
                    higher->next = head;
                    higher = higher->next;
                }
            }
            head = head->next;
        }
        if (lowHead == nullptr)
            return highHead;
        lower->next = highHead;
        if (higher != nullptr)
            higher->next = nullptr;
        return lowHead;
    }
    ListNode *partition_waste_memory(ListNode *head, int x)
    {
        if (head == nullptr)
            return nullptr;
        ListNode *lowHead = nullptr, *highHead = nullptr;
        ListNode *lower = nullptr, *higher = nullptr;
        while (head != nullptr)
        {
            if (head->val < x)
            {
                if (lower == nullptr)
                {
                    lower = new ListNode(head->val);
                    lowHead = lower;
                }
                else
                {
                    lower->next = new ListNode(head->val);
                    lower = lower->next;
                }
            }
            else
            {
                if (higher == nullptr)
                {
                    higher = new ListNode(head->val);
                    highHead = higher;
                }
                else
                {
                    higher->next = new ListNode(head->val);
                    higher = higher->next;
                }
            }
            head = head->next;
        }
        if (lowHead == nullptr)
            return highHead;
        lower->next = highHead;
        return lowHead;
    }
};

int main()
{
    int arr[] = {1};
    int x = 2;
    ListNode *head = CreateLinkedList(arr, sizeof(arr) / sizeof(arr[0]));
    Solution s;
    PrintLinkedList(head);
    ListNode *result = s.partition(head, x);
    PrintLinkedList(result);

    // DeleteLinkedList(head);
    DeleteLinkedList(result);
    return 0;
}